import os
import re
from fastapi import APIRouter, WebSocket
from loguru import logger
from datetime import datetime
from typing import Generator, Any
from app.graph.llm import deepseek_chat

from .asr import do_funasr
from .tts import text_to_speech

router = APIRouter(prefix="/ws", tags=["web socket"])


async def voice_response(text: str, websocket: WebSocket):
    # TODO: websocket 为正常关闭似乎会挂
    tokens = deepseek_chat(
        [
            {
                "role": "user",
                "content": f"你是一个语音助手。你在回答问题时，回答的内容中，只用中文的逗号和中文的句号断句，以便语音机器人能够正确读取你回答的文本消息。以下式用户的问题:\n\n\n{text}",
            }
        ]
    )
    # ④ deepseek tokens

    sentences = tokens_to_sentences(tokens)

    for s in sentences:
        b = await text_to_speech(s)
        if b:
            logger.debug("    [WS] send bytes")
            await websocket.send_bytes(b)
            await websocket.send_text(s)

    await websocket.close(code=1000, reason="All done!")
    logger.debug("[ws] ALL DONE!")


@router.websocket("/chat-text")
async def chat_text(websocket: WebSocket):
    await websocket.accept()
    text = await websocket.receive_text()
    await voice_response(text, websocket)


@router.websocket("/chat")
async def chat(websocket: WebSocket):
    await websocket.accept()
    file = await websocket.receive_bytes()

    # 不能直接把 ogg/opus 字节串直接作为 input 传给 funasr，因为这样做效果极差 😂
    filename = f"temp-{datetime.now().strftime('%Y%m%d%H%M%S%f')}.ogg"
    with open(filename, "wb") as tmp:
        tmp.write(file)  # write ogg bytes into temp file

    text = do_funasr(filename)  # funasr
    logger.debug("asr text: {}", text)

    os.remove(filename)  # clean temp file

    await voice_response(text, websocket)


def tokens_to_sentences(tokens: Generator[str, Any, None]):
    buffer = ""
    for token in tokens:
        buffer += token
        result = re.search(r"[。！？，\n]+", buffer)
        if result:  # 搜到了句子末尾
            (s, e) = result.span()
            sentence = buffer[0 : (s + 1)]  # 前面的部分就是句子
            buffer = buffer[e:]  # 剩下的部分继续留在 buffer 里
            if sentence:
                logger.debug("  [Sentence] yield sentence : {}", sentence)
                yield sentence
    if buffer:
        yield buffer
